# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import division
import pygame
from pygame import Rect
import random
import numpy

dirs = (1, 1j, -1, -1j)

INF = float('Inf')

WHITE = pygame.Color('white')
BLACK = pygame.Color('black')

class Gameplay (object):
    SCALE = 30

    def __init__ (self):
        self.level = 1
        self.score = 0
        self.lines = 0

    def piece_size (self):
        return self.SCALE

    def falling_piece_color (self, omino):
        """Return a random bright color."""
        color = [0xff, random.randint(0,0xff), random.randint(0,0xff)]
        random.shuffle(color)
        return tuple(color)

    def cemented_piece_color (self, omino):
        """The color of the piece after it is cemented.

        Must not return black.
        """

        return omino.color

    def snappiness (self):
        return self.SCALE

    def cement_time (self):
        return self.SCALE

    def level_up (self):
        self.level += 1

    def omino_size (self):
        return 4

    def field_size (self):
        return (10*self.SCALE, 20*self.SCALE)

    def gap_size (self):
        return 0

    def erase_color (self):
        return (0, 0, 0)

    def erase_time (self):
        return 20

    def got_rows (self, rows):
        """Account for rows of pixels completed."""
        self.got_lines(rows/self.SCALE)

    def got_lines (self, lines):
        """Score actual complete lines (a line is some number of rows)."""
        if lines:
            self.score += self.level*int(100 * 2**(lines-1))
            self.lines += lines
            if self.lines >= self.level*5:
                self.level_up()

class MisGameplay (Gameplay):
    TOT_GAP_TOLERANCE = 0.99
    def __init__ (self):
        Gameplay.__init__(self)
        self.factor = self.SCALE

    def piece_size (self):
        if self.factor == 0:
            return random.randint(1, self.SCALE)
        else:
            return random.choice([self.factor, self.SCALE])

    def snappiness (self):
        return max(self.factor, 1)

    def level_up (self):
        Gameplay.level_up(self)
        for i in range(self.factor-1, 0, -1):
            if self.SCALE % i == 0:
                self.factor = i
                return
        self.factor = 0

    def gap_size (self):
        return self.TOT_GAP_TOLERANCE*self.SCALE

    def got_rows (self, rows):
        Gameplay.got_lines(self, 2 * rows/(self.SCALE + self.snappiness()))

class MultiGameplay (Gameplay):
    def omino_size (self):
        return random.randrange(max(1,5-self.level), self.level+4)

    def field_size (self):
        return (15*self.SCALE, 20*self.SCALE)

class Omino (object):
    def __init__ (self, gameplay):
        self.gameplay = gameplay
        self.scale = gameplay.piece_size()
        self.gravity = 1
        self.horz_speed = 1
        self.ttl = gameplay.cement_time()

        self.moving_state = 0
        self.snappiness = gameplay.snappiness()

        self.field = None
        self.rect = Rect(0,0,0,0)
        self.build(gameplay.omino_size())
        self.color = gameplay.falling_piece_color(self)

    def set_field (self, field):
        self.field = field
        self.rect.centerx = field.rect.width // 2
        self.rect.left -= self.rect.left % self.snappiness
        self.rect.bottom = 0

    def build (self, n):
        assert n >= 1
        self.blocks = [0j]

        while len(self.blocks) < n:
            parent = random.choice(self.blocks)
            dir = random.choice(dirs)
            child = parent + dir
            if child not in self.blocks:
                self.blocks.append(child)

        self.normalize()

    def normalize (self):
        left = top = right = bottom = 0
        for block in self.blocks:
            left = min(left, block.real)
            right = max(right, block.real)
            top = min(top, block.imag)
            bottom = max(bottom, block.imag)

        self.blocks = [block-(left+top*1j)
                        for block in self.blocks]
        self.rects = [Rect(self.scale*block.real, self.scale*block.imag,
                            self.scale, self.scale)
                        for block in self.blocks]

        oldcenter = self.rect.center
        self.rect.size = self.rects[0].unionall(self.rects).size
        self.rect.center = oldcenter

        if self.rect.width//self.scale % 2 == 0:
            self.rect.left = -(-self.rect.left // self.snappiness) * self.snappiness
        else:
            self.rect.left = (self.rect.left // self.snappiness) * self.snappiness

        self.clamp()

    def clamp (self):
        if self.field is None:
            self.rect.left = 0
            self.rect.top = 0
        else:
            self.rect.left = max(self.rect.left, 0)
            self.rect.right = min(self.rect.right, self.field.rect.width)

    def rotate (self):
        old = self.blocks[:], self.rects[:], Rect(self.rect)
        self.blocks = [block*1j for block in self.blocks]
        self.normalize()

        if self.collides():
            self.blocks, self.rects, self.rect = old

    def left (self):
        self.moving_state = -1

    def unleft (self):
        if self.moving_state == -1:
            self.moving_state = 0

            old = Rect(self.rect)
            self.rect.left = (self.rect.left // self.snappiness) * self.snappiness
            if self.collides():
                self.rect = old

    def right (self):
        self.moving_state = 1

    def unright (self):
        if self.moving_state == 1:
            self.moving_state = 0

            old = Rect(self.rect)
            self.rect.left = -(-self.rect.left // self.snappiness) * self.snappiness
            if self.collides():
                self.rect = old

    def update (self):
        if self.has_landed():
            self.ttl -= 1
        else:
            self.ttl = self.gameplay.cement_time()
            self.rect.move_ip(0, self.gravity)

        if self.moving_state:
            if ((self.rect.left == 0
                    and self.moving_state < 0) or
                (self.rect.right == self.field.rect.width
                    and self.moving_state > 0)):
                return

            pixels = pygame.surfarray.pixels2d(self.field.image)
            for rect in self.rects:
                rect = rect.move(*self.rect.topleft)
                if rect.bottom <= 0:
                    continue

                if self.moving_state < 0:
                    x = rect.left + self.moving_state
                else:
                    x = rect.right + self.moving_state - 1
                y = max(rect.top, 0)
                if numpy.any(pixels[x, y:rect.bottom]):
                    break
            else:
                self.rect.move_ip(self.moving_state*self.horz_speed, 0)
                self.clamp()

    def draw (self, surf, into=None, color=None):
        if color is None:
            color = self.color
        for rect in self.rects:
            rect = rect.move(*self.rect.topleft)
            if into:
                rect.move_ip(into.topleft)
                rect = rect.clip(into)
            surf.fill(color, rect)

    def drop (self):
        drop_dist = self.field.rect.height - self.rect.bottom

        for rect in self.rects:
            dy = rect.height
            rect = rect.move(self.rect.topleft)
            y = max(rect.bottom, 0)
            shadow = pygame.surfarray.pixels2d(self.field.image)
            shadow = shadow[rect.left:rect.right, y:]
            rows = numpy.any(shadow, axis=0)

            try:
                drop_dist = min(drop_dist, numpy.nonzero(rows)[0][0])
            except IndexError:
                pass

        self.rect.bottom += drop_dist

    def cement (self):
        self.color = self.gameplay.cemented_piece_color(self)
        return self.field.cement(self)

    def has_landed (self):
        if self.rect.bottom <= 0:
            return False
        elif self.rect.bottom == self.field.rect.height:
            return 'bottom'

        pixels = pygame.surfarray.pixels2d(self.field.image)
        for rect in self.rects:
            rect = rect.move(*self.rect.topleft)
            if rect.bottom <= 0:
                continue
            if numpy.any(pixels[rect.left:rect.right, rect.bottom]):
                return 'block'

        return False

    def collides (self):
        pixels = pygame.surfarray.pixels2d(self.field.image)
        for rect in self.rects:
            rect = rect.move(*self.rect.topleft)
            if rect.bottom <= 0:
                continue
            if numpy.any(pixels[rect.left:rect.right,
                                max(rect.top, 0):rect.bottom]):
                return True
        else:
            return False

class Field (object):
    def __init__ (self, pos, gameplay):
        self.gameplay = gameplay
        size = gameplay.field_size()
        self.rect = Rect(pos, size)
        self.image = pygame.Surface(size).convert()
        self.lines_to_erase = set()

    def cement (self, omino):
        omino.draw(self.image)
        pixels = pygame.surfarray.pixels2d(self.image)
        for y in range(omino.rect.bottom-1, omino.rect.top-1, -1):
            if self.is_line(pixels, y):
                self.lines_to_erase.add(y)
                self.ttl = self.gameplay.erase_time()

    def erase (self, omino):
        lines = 0
        pixels = pygame.surfarray.pixels2d(self.image)
        for y in range(omino.rect.bottom-1, omino.rect.top-1, -1):
            if y in self.lines_to_erase:
                lines += 1
            else:
                pixels[:, y+lines] = pixels[:, y]

        self.lines_to_erase.clear()

        pixels[:, lines:omino.rect.top+lines] = pixels[:, :omino.rect.top]
        pixels[:, :lines] = 0

        return lines

    def draw (self, surf):
        surf.blit(self.image, self.rect.topleft)
        for line in self.lines_to_erase:
            rect = Rect(self.rect.left, line+self.rect.top, self.rect.width, 1)
            surf.fill(self.gameplay.erase_color(), rect)

    def is_line (self, pixels, y):
        return (numpy.sum(pixels[:, y] == 0) <= self.gameplay.gap_size())

def draw_text (font, text, surf, pos, backcolor=BLACK, forecolor=WHITE):
    surf.blit(font.render(text, True, forecolor, backcolor), pos)

def main ():
    pygame.init()
    background = pygame.image.load('background.png')
    screen = pygame.display.set_mode(background.get_size())
    background = background.convert()
    screen.blit(background, (0, 0))
    offset = (90, 84)

    gameplay = MisGameplay()

    field = Field(offset, gameplay)
    font = pygame.font.Font(None, 50)

    clock = pygame.time.Clock()
    paused = False

    omino = Omino(gameplay)
    omino.set_field(field)
    next = Omino(gameplay)
    old = None
    while True:
        clock.tick(100)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return
            elif event.type == pygame.KEYDOWN:
                # TODO: make keys configurable
                keys = {pygame.K_i: omino.rotate,
                        pygame.K_j: omino.left,
                        pygame.K_l: omino.right,
                        pygame.K_k: omino.drop}
                if not paused and event.key in keys:
                    keys[event.key]()
                elif event.key == pygame.K_SPACE:
                    paused = not paused
                elif event.key == pygame.K_ESCAPE:
                    return
            elif event.type == pygame.KEYUP:
                keys = {pygame.K_j: omino.unleft,
                        pygame.K_l: omino.unright}
                if not paused and event.key in keys:
                    keys[event.key]()

        if not paused:
            if field.lines_to_erase:
                field.ttl -= 1
                if field.ttl == 0:
                    rows = field.erase(old)
                    gameplay.got_rows(rows)
            else:
                omino.update()

        if not paused and omino.ttl == 0 and omino.has_landed():
            # remember which keys are down
            next.moving_state = omino.moving_state
            # snap if moving
            omino.unleft()
            omino.unright()
            if omino.has_landed():
                omino.cement()
                old = omino
                omino = next
                omino.set_field(field)
                next = Omino(gameplay)

        field.draw(screen)
        omino.draw(screen, field.rect)

        screen.fill(BLACK, Rect(570,370,180,180))
        next.draw(screen, Rect(600, 400, 1000, 1000))
        draw_text(font, 'Level %i' % gameplay.level, screen, (600, 100))
        draw_text(font, '%i points' % gameplay.score, screen, (600, 300))

        pygame.display.flip()

if __name__ == '__main__':
    main()
